`timescale 1ns / 1ps

`define SLICE(BIT, IDX) (``BIT``) * (``IDX``) + (``BIT``) - 1 : (``BIT``) * (``IDX``)

module max2pool
#(
    parameter N_IO      = 4,
    parameter N_CH      = 48,
    parameter N_COL     = 320,
    parameter BIT       = 8,
    parameter RAM_TYPE  = "block",
    parameter RAM_LATENCY = 2
)
(
    input   clk,
    input   rst,
    
    input   i_vld,
    output  i_rdy,
    input   [N_IO*BIT-1 : 0]      i_data,
    
    output  o_vld,
    input   o_rdy,
    output  [N_IO*BIT-1 : 0]      o_data
);

localparam FOLD = N_CH / N_IO;
localparam D_WIDTH = N_IO * BIT;
localparam D_DEPTH = N_COL / 2 * FOLD;
localparam A_WIDTH = $clog2(D_DEPTH);
localparam [A_WIDTH-1: 0] A_ZERO = {A_WIDTH{1'b0}};
localparam [A_WIDTH-1: 0] A_MAX = D_DEPTH - 1;

wire [1: 0]             cstate;

wire                    enable;
reg                     row;    // row_lsb
reg                     col;    // col_lsb
wire                    fod_last;
reg  [A_WIDTH-1 : 0]    addr;
wire                    addr_last;

reg                     w_ena;
reg  [A_WIDTH-1 : 0]    w_addr;

wire                    r_ena;
wire                    r_ena_cond;
reg                     r_ena_pipe;     // RAM_LATENCY == 2
wire [A_WIDTH-1 : 0]    r_addr;
wire [D_WIDTH-1 : 0]    r_data;

wire                    m_sel;
reg                     m_vld;
reg  [D_WIDTH-1 : 0]    m_data;

wire [1: 0]             ystate;
wire                    y_vld;
wire [A_WIDTH-1 : 0]    y_addr;
wire [D_WIDTH-1 : 0]    y_data;

wire [D_WIDTH-1 : 0]    mx_out;

assign i_rdy = enable;
assign o_vld = m_vld;
assign o_data = m_data;

assign cstate = { row, col };
assign enable = ~m_vld | o_rdy;
assign r_ena  = enable & (r_ena_cond | r_ena_pipe);
assign r_ena_cond = i_vld & (cstate != 2'b00);
assign r_addr = addr;
assign m_sel  = ystate != 2'b00;
assign addr_last = addr == A_MAX;

always @(posedge clk)
begin
    if (rst) begin
        row <= 1'b0;
        col <= 1'b0;
        addr <= A_ZERO;
    end
    else begin
        if (enable & i_vld) begin
            row <= row ^ (fod_last & col & addr_last);
            col <= col ^  fod_last;
            if (fod_last & col & addr_last)
                addr <= A_ZERO;
            else if (fod_last & ~col)
                addr <= addr - (FOLD - 1);
            else
                addr <= addr + 1'b1;
        end
    end
end

always @(posedge clk)
begin
    if (rst) begin
        m_vld <= 1'b0;
        w_ena <= 1'b0;
        r_ena_pipe <= 1'b0;
    end
    else begin
        if (enable) begin
            m_vld <= y_vld & (ystate == 2'b11);
            w_ena <= y_vld & (ystate != 2'b11);
            r_ena_pipe <= r_ena_cond;
        end
        if (enable & y_vld)
            m_data <= m_sel ? mx_out : y_data;
        if (enable & y_vld & (ystate != 2'b11))
            w_addr <= y_addr;
    end
end

zq_counter #(
    .N  ( FOLD )
) inst_cnt_fold
(
    .clk    (clk),
    .rst    (rst),
    .clken  (enable & i_vld),
    .last   (fod_last),
    .out    ()
);

shift_reg_rst #(
    .N     ( RAM_LATENCY ),
    .WIDTH ( 2 + 1       )
)
inst_ctrl_line (
    .clk                     ( clk      ),
    .rst                     ( rst      ),

    .clken                   ( enable   ),
    .i_data                  ({ cstate, i_vld }),
    .o_data                  ({ ystate, y_vld })
);

shift_reg #(
    .N     ( RAM_LATENCY        ),
    .WIDTH ( A_WIDTH + D_WIDTH  )
)
inst_delay_line (
    .clk                     ( clk      ),
    .clken                   ( enable   ),

    .i_data                  ({   addr, i_data }),
    .o_data                  ({ y_addr, y_data })
);

zq_max #(
    .N_IO   ( N_IO  ),
    .BIT    ( BIT   )
)
inst_zqmax (
    .ain    ( y_data ),
    .bin    ( r_data ),
    .out    ( mx_out )
);

zq_sdpram #(
    .ADDR_WIDTH ( A_WIDTH       ),
    .DATA_WIDTH ( D_WIDTH       ),
    .DEPTH      ( D_DEPTH       ),
    .LATENCY    ( RAM_LATENCY   ),
    .RAMTYPE    ( RAM_TYPE      )
)
inst_sdpram (
    .clk                     ( clk      ),

    .ena_w                   ( enable   ),
    .wea                     ( w_ena    ),
    .addr_w                  ( w_addr   ),
    .din                     ( m_data   ),

    .ena_r                   ( r_ena    ),
    .addr_r                  ( r_addr   ),
    .dout                    ( r_data   )
);

endmodule


module zq_max
#(
    parameter N_IO      = 4,
    parameter BIT       = 8
)
(
    input   [N_IO*BIT-1 : 0]      ain,
    input   [N_IO*BIT-1 : 0]      bin,
    
    output  [N_IO*BIT-1 : 0]      out
);

wire [BIT-1 : 0] x [N_IO-1 : 0];
wire [BIT-1 : 0] y [N_IO-1 : 0];
wire [BIT-1 : 0] z [N_IO-1 : 0];

genvar i;
generate
    for (i = 0; i < N_IO; i = i + 1)
    begin
        assign x[i] = ain[ `SLICE(BIT, i) ];
        assign y[i] = bin[ `SLICE(BIT, i) ];
        assign z[i] = x[i] > y[i] ? x[i] : y[i];
        assign out[ `SLICE(BIT, i) ] = z[i];
    end
endgenerate

endmodule
